"""
TensorBoard可视化模块，实现数据可视化功能。
"""
import os
from enum import Enum
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass
import torch
from torch.utils.tensorboard import SummaryWriter

class MetricType(Enum):
    """指标类型枚举"""
    GRADIENT = "gradient"
    ACTIVATION = "activation"
    PARAMETER = "parameter"
    OPTIMIZER = "optimizer"
    ANOMALY = "anomaly"
    PERFORMANCE = "performance"
    MEMORY = "memory"
    CUSTOM = "custom"

@dataclass
class MetricData:
    """指标数据类"""
    name: str
    value: Union[float, Dict[str, float], torch.Tensor]
    step: int
    metric_type: MetricType
    metadata: Optional[Dict[str, Any]] = None

class TensorBoardVisualizer:
    """TensorBoard可视化器"""
    
    def __init__(self, log_dir: str):
        """
        初始化TensorBoard可视化器
        
        Args:
            log_dir: 日志目录
        """
        os.makedirs(log_dir, exist_ok=True)
        self.writer = SummaryWriter(log_dir)
    
    def add_scalar(self, 
                 name: str, 
                 value: float, 
                 step: int, 
                 metric_type: MetricType) -> None:
        """
        添加标量指标
        
        Args:
            name: 指标名称
            value: 指标值
            step: 步数
            metric_type: 指标类型
        """
        # 构建标签，包含指标类型
        tag = f"{metric_type.value}/{name}"
        self.writer.add_scalar(tag, value, step)
    
    def add_scalars(self, 
                  name: str, 
                  value_dict: Dict[str, float], 
                  step: int, 
                  metric_type: MetricType) -> None:
        """
        添加多个标量指标
        
        Args:
            name: 指标名称前缀
            value_dict: 值字典
            step: 步数
            metric_type: 指标类型
        """
        # 构建标签，包含指标类型
        tag = f"{metric_type.value}/{name}"
        self.writer.add_scalars(tag, value_dict, step)
    
    def add_histogram(self, 
                    name: str, 
                    values: torch.Tensor, 
                    step: int, 
                    metric_type: MetricType) -> None:
        """
        添加直方图
        
        Args:
            name: 指标名称
            values: 张量值
            step: 步数
            metric_type: 指标类型
        """
        # 构建标签，包含指标类型
        tag = f"{metric_type.value}/{name}"
        self.writer.add_histogram(tag, values, step)
    
    def add_text(self, 
               name: str, 
               text: str, 
               step: int, 
               metric_type: MetricType) -> None:
        """
        添加文本
        
        Args:
            name: 指标名称
            text: 文本内容
            step: 步数
            metric_type: 指标类型
        """
        # 构建标签，包含指标类型
        tag = f"{metric_type.value}/{name}"
        self.writer.add_text(tag, text, step)
    
    def add_metric(self, metric: MetricData) -> None:
        """
        添加指标数据
        
        Args:
            metric: 指标数据对象
        """
        if isinstance(metric.value, float) or isinstance(metric.value, int):
            self.add_scalar(metric.name, float(metric.value), metric.step, metric.metric_type)
        elif isinstance(metric.value, dict):
            self.add_scalars(metric.name, metric.value, metric.step, metric.metric_type)
        elif isinstance(metric.value, torch.Tensor):
            self.add_histogram(metric.name, metric.value, metric.step, metric.metric_type)
        elif isinstance(metric.value, str):
            self.add_text(metric.name, metric.value, metric.step, metric.metric_type)
    
    def add_metrics(self, metrics: List[MetricData]) -> None:
        """
        添加多个指标数据
        
        Args:
            metrics: 指标数据列表
        """
        for metric in metrics:
            self.add_metric(metric)
    
    def add_model_graph(self, model: torch.nn.Module, input_tensor: torch.Tensor) -> None:
        """
        添加模型图
        
        Args:
            model: 模型
            input_tensor: 输入张量
        """
        self.writer.add_graph(model, input_tensor)
    
    def close(self) -> None:
        """关闭写入器"""
        self.writer.close()
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.close() 